#include "MeshExtender.h"

MeshExtender::MeshExtender(std::vector<Point3D>& nodes, 
                         std::vector<Tetrahedron>& elements,
                         const ExtendParams& params,
                         const GridSizes& gridSizes)
    : nodes_(nodes), elements_(elements), params_(params), gridSizes_(gridSizes) {
    
    // 初始化储层范围
    x_min_ = nodes[0].x;
    x_max_ = nodes[0].x;
    y_min_ = nodes[0].y;
    y_max_ = nodes[0].y;
    z_min_ = nodes[0].z;
    z_max_ = nodes[0].z;
    
    for (const auto& node : nodes) {
        x_min_ = std::min(x_min_, node.x);
        x_max_ = std::max(x_max_, node.x);
        y_min_ = std::min(y_min_, node.y);
        y_max_ = std::max(y_max_, node.y);
        z_min_ = std::min(z_min_, node.z);
        z_max_ = std::max(z_max_, node.z);
    }
    
    // 初始化cellZones_，所有原始单元标记为1
    cellZones_.resize(elements.size(), 1);
    
    // 计算延拓尺寸
    computeExtendSizes();
}

void MeshExtender::computeExtendSizes() {
    // 计算储层范围尺寸
    delta_reservoir_x_ = x_max_ - x_min_;
    delta_reservoir_y_ = y_max_ - y_min_;
    delta_reservoir_z_ = z_max_ - z_min_;
    
    // 使用GridGenerator计算的网格尺寸
    extend_size_x_ = (params_.ratioX * delta_reservoir_x_) / params_.numExtendX;
    extend_size_x_middle_ = gridSizes_.dx;  // 使用传入的dx
    extend_size_y_ = (params_.ratioY * delta_reservoir_y_) / params_.numExtendY;
    extend_size_y_middle_ = gridSizes_.dy;  // 使用传入的dy
    
    // 修改z方向延拓尺寸的计算
    double model_z_max = -params_.downElevation;
    double model_z_min = -params_.upElevation;
    
    extend_size_z_up_ = std::abs(z_max_ - model_z_max) / params_.numExtendUp;
    extend_size_z_down_ = std::abs(z_min_ - model_z_min) / params_.numExtendDown;
    extend_size_z_middle_ = gridSizes_.dz;  // 使用传入的dz
}

void MeshExtender::extend() {
    // 按顺序生成各个区域的网格
    generateSurroundingRock();
    generateOverlyingRock();
    generateOverlyingSurroundingRock();
    generateUnderlyingRock();
    generateUnderlyingSurroundingRock();
}

void MeshExtender::generateHexMesh(
    double x_start, double x_end, double x_size,
    double y_start, double y_end, double y_size,
    double z_start, double z_end, double z_size,
    int zone_type) {
    
    // 计算各方向的网格数量
    int nx = std::ceil((x_end - x_start) / x_size);
    int ny = std::ceil((y_end - y_start) / y_size);
    int nz = std::ceil((z_end - z_start) / z_size);
    
    // 调整网格尺寸以保证精确覆盖区域
    x_size = (x_end - x_start) / nx;
    y_size = (y_end - y_start) / ny;
    z_size = (z_end - z_start) / nz;
    
    // 生成节点
    size_t base_node_index = nodes_.size();
    for (int k = 0; k <= nz; ++k) {
        for (int j = 0; j <= ny; ++j) {
            for (int i = 0; i <= nx; ++i) {
                double x = x_start + i * x_size;
                double y = y_start + j * y_size;
                double z = z_start + k * z_size;
                nodes_.emplace_back(x, y, z);
            }
        }
    }
    
    // 生成六面体并剖分为四面体
    int nx1 = nx + 1;
    int ny1 = ny + 1;
    for (int k = 0; k < nz; ++k) {
        for (int j = 0; j < ny; ++j) {
            for (int i = 0; i < nx; ++i) {
                // 构建六面体的8个顶点索引
                std::vector<size_t> hexNodes = {
                    base_node_index + k * nx1 * ny1 + j * nx1 + i,
                    base_node_index + k * nx1 * ny1 + j * nx1 + (i + 1),
                    base_node_index + k * nx1 * ny1 + (j + 1) * nx1 + (i + 1),
                    base_node_index + k * nx1 * ny1 + (j + 1) * nx1 + i,
                    base_node_index + (k + 1) * nx1 * ny1 + j * nx1 + i,
                    base_node_index + (k + 1) * nx1 * ny1 + j * nx1 + (i + 1),
                    base_node_index + (k + 1) * nx1 * ny1 + (j + 1) * nx1 + (i + 1),
                    base_node_index + (k + 1) * nx1 * ny1 + (j + 1) * nx1 + i
                };
                
                // 将六面体剖分为四面体
                auto tets = splitHexToTets(hexNodes);
                
                // 添加四面体并设置区域标记
                for (const auto& tet : tets) {
                    elements_.push_back(tet);
                    cellZones_.push_back(zone_type);
                }
            }
        }
    }
}

std::vector<Tetrahedron> MeshExtender::splitHexToTets(const std::vector<size_t>& hexNodes) {
    // 六面体的8个顶点按照以下顺序：
    // hexNodes[0] = v000 (左下后)
    // hexNodes[1] = v100 (右下后)
    // hexNodes[2] = v110 (右上后)
    // hexNodes[3] = v010 (左上后)
    // hexNodes[4] = v001 (左下前)
    // hexNodes[5] = v101 (右下前)
    // hexNodes[6] = v111 (右上前)
    // hexNodes[7] = v011 (左上前)

    // 使用Kuhn剖分方案，确保四面体的正确方向
    return {
        // 每个四面体的顶点顺序遵循右手法则
        Tetrahedron{hexNodes[0], hexNodes[2], hexNodes[7], hexNodes[3]},  // v000, v110, v011, v010
        Tetrahedron{hexNodes[0], hexNodes[5], hexNodes[2], hexNodes[1]},  // v000, v101, v110, v100
        Tetrahedron{hexNodes[0], hexNodes[7], hexNodes[5], hexNodes[4]},  // v000, v011, v101, v001
        Tetrahedron{hexNodes[2], hexNodes[5], hexNodes[7], hexNodes[6]},  // v110, v101, v011, v111
        Tetrahedron{hexNodes[0], hexNodes[2], hexNodes[5], hexNodes[7]}   // v000, v110, v101, v011
    };
}

void MeshExtender::generateSurroundingRock() {
    // 生成围岩区域网格
    // 前后
    generateHexMesh(x_min_, x_max_, extend_size_x_middle_,
                   y_min_ - params_.ratioY * delta_reservoir_y_, y_min_, extend_size_y_,
                   z_min_, z_max_, extend_size_z_middle_,
                   2);
    generateHexMesh(x_min_, x_max_, extend_size_x_middle_,
                   y_max_, y_max_ + params_.ratioY * delta_reservoir_y_, extend_size_y_,
                   z_min_, z_max_, extend_size_z_middle_,
                   2);
    
    // 左右
    generateHexMesh(x_min_ - params_.ratioX * delta_reservoir_x_, x_min_, extend_size_x_,
                   y_min_, y_max_, extend_size_y_middle_,
                   z_min_, z_max_, extend_size_z_middle_,
                   2);
    generateHexMesh(x_max_, x_max_ + params_.ratioX * delta_reservoir_x_, extend_size_x_,
                   y_min_, y_max_, extend_size_y_middle_,
                   z_min_, z_max_, extend_size_z_middle_,
                   2);
    
    // 四个角
    generateHexMesh(x_min_ - params_.ratioX * delta_reservoir_x_, x_min_, extend_size_x_,
                   y_min_ - params_.ratioY * delta_reservoir_y_, y_min_, extend_size_y_,
                   z_min_, z_max_, extend_size_z_middle_,
                   2);
    generateHexMesh(x_max_, x_max_ + params_.ratioX * delta_reservoir_x_, extend_size_x_,
                   y_min_ - params_.ratioY * delta_reservoir_y_, y_min_, extend_size_y_,
                   z_min_, z_max_, extend_size_z_middle_,
                   2);
    generateHexMesh(x_min_ - params_.ratioX * delta_reservoir_x_, x_min_, extend_size_x_,
                   y_max_, y_max_ + params_.ratioY * delta_reservoir_y_, extend_size_y_,
                   z_min_, z_max_, extend_size_z_middle_,
                   2);
    generateHexMesh(x_max_, x_max_ + params_.ratioX * delta_reservoir_x_, extend_size_x_,
                   y_max_, y_max_ + params_.ratioY * delta_reservoir_y_, extend_size_y_,
                   z_min_, z_max_, extend_size_z_middle_,
                   2);
}

void MeshExtender::generateOverlyingRock() {
    // 生成上覆岩石区域网格（区域3）
    double model_z_max = -params_.downElevation;  // 正确的上边界
    generateHexMesh(
        x_min_, x_max_, extend_size_x_middle_,
        y_min_, y_max_, extend_size_y_middle_,
        z_max_, model_z_max, extend_size_z_up_,  // 从小到大
        3
    );
}

void MeshExtender::generateOverlyingSurroundingRock() {
    // 生成上覆围岩区域网格（区域4）
    double model_z_max = -params_.downElevation;  // 正确的上边界
    
    // 前后
    generateHexMesh(x_min_, x_max_, extend_size_x_middle_,
                   y_min_ - params_.ratioY * delta_reservoir_y_, y_min_, extend_size_y_,
                   z_max_, model_z_max, extend_size_z_up_,  // 从小到大
                   4);
    generateHexMesh(x_min_, x_max_, extend_size_x_middle_,
                   y_max_, y_max_ + params_.ratioY * delta_reservoir_y_, extend_size_y_,
                   z_max_, model_z_max, extend_size_z_up_,
                   4);
    
    // 左右
    generateHexMesh(x_min_ - params_.ratioX * delta_reservoir_x_, x_min_, extend_size_x_,
                   y_min_, y_max_, extend_size_y_middle_,
                   z_max_, model_z_max, extend_size_z_up_,
                   4);
    generateHexMesh(x_max_, x_max_ + params_.ratioX * delta_reservoir_x_, extend_size_x_,
                   y_min_, y_max_, extend_size_y_middle_,
                   z_max_, model_z_max, extend_size_z_up_,
                   4);
    
    // 四个角
    generateHexMesh(x_min_ - params_.ratioX * delta_reservoir_x_, x_min_, extend_size_x_,
                   y_min_ - params_.ratioY * delta_reservoir_y_, y_min_, extend_size_y_,
                   z_max_, model_z_max, extend_size_z_up_,
                   4);
    generateHexMesh(x_max_, x_max_ + params_.ratioX * delta_reservoir_x_, extend_size_x_,
                   y_min_ - params_.ratioY * delta_reservoir_y_, y_min_, extend_size_y_,
                   z_max_, model_z_max, extend_size_z_up_,
                   4);
    generateHexMesh(x_min_ - params_.ratioX * delta_reservoir_x_, x_min_, extend_size_x_,
                   y_max_, y_max_ + params_.ratioY * delta_reservoir_y_, extend_size_y_,
                   z_max_, model_z_max, extend_size_z_up_,
                   4);
    generateHexMesh(x_max_, x_max_ + params_.ratioX * delta_reservoir_x_, extend_size_x_,
                   y_max_, y_max_ + params_.ratioY * delta_reservoir_y_, extend_size_y_,
                   z_max_, model_z_max, extend_size_z_up_,
                   4);
}

void MeshExtender::generateUnderlyingRock() {
    // 生成下覆岩石区域网格（区域5）
    double model_z_min = -params_.upElevation;  // 正确的下边界
    generateHexMesh(
        x_min_, x_max_, extend_size_x_middle_,
        y_min_, y_max_, extend_size_y_middle_,
        model_z_min, z_min_, extend_size_z_down_,  // 从小到大
        5
    );
}

void MeshExtender::generateUnderlyingSurroundingRock() {
    // 生成下覆围岩区域网格（区域6）
    double model_z_min = -params_.upElevation;  // 正确的下边界
    
    // 前后
    generateHexMesh(x_min_, x_max_, extend_size_x_middle_,
                   y_min_ - params_.ratioY * delta_reservoir_y_, y_min_, extend_size_y_,
                   model_z_min, z_min_, extend_size_z_down_,  // 从小到大
                   6);
    generateHexMesh(x_min_, x_max_, extend_size_x_middle_,
                   y_max_, y_max_ + params_.ratioY * delta_reservoir_y_, extend_size_y_,
                   model_z_min, z_min_, extend_size_z_down_,
                   6);
    
    // 左右
    generateHexMesh(x_min_ - params_.ratioX * delta_reservoir_x_, x_min_, extend_size_x_,
                   y_min_, y_max_, extend_size_y_middle_,
                   model_z_min, z_min_, extend_size_z_down_,
                   6);
    generateHexMesh(x_max_, x_max_ + params_.ratioX * delta_reservoir_x_, extend_size_x_,
                   y_min_, y_max_, extend_size_y_middle_,
                   model_z_min, z_min_, extend_size_z_down_,
                   6);
    
    // 四个角
    generateHexMesh(x_min_ - params_.ratioX * delta_reservoir_x_, x_min_, extend_size_x_,
                   y_min_ - params_.ratioY * delta_reservoir_y_, y_min_, extend_size_y_,
                   model_z_min, z_min_, extend_size_z_down_,
                   6);
    generateHexMesh(x_max_, x_max_ + params_.ratioX * delta_reservoir_x_, extend_size_x_,
                   y_min_ - params_.ratioY * delta_reservoir_y_, y_min_, extend_size_y_,
                   model_z_min, z_min_, extend_size_z_down_,
                   6);
    generateHexMesh(x_min_ - params_.ratioX * delta_reservoir_x_, x_min_, extend_size_x_,
                   y_max_, y_max_ + params_.ratioY * delta_reservoir_y_, extend_size_y_,
                   model_z_min, z_min_, extend_size_z_down_,
                   6);
    generateHexMesh(x_max_, x_max_ + params_.ratioX * delta_reservoir_x_, extend_size_x_,
                   y_max_, y_max_ + params_.ratioY * delta_reservoir_y_, extend_size_y_,
                   model_z_min, z_min_, extend_size_z_down_,
                   6);
} 